import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch.utils.data.sampler import Sampler
import numpy as np
import torch
import random
import os
import math
import time

class ShuffleSampler(Sampler):

    def __init__(self, data, seed):

        lenth = int(len(data))
        self.indices = list(
            [index for index in range(lenth)]
        )
        random.seed(seed)
        random.shuffle(self.indices)
    
    def __iter__(self, ):

        return iter(self.indices)
    
    def __len__(self, ):

        return len(self.indices)

class ExculdeClassSampler(Sampler):

    def __init__(self, data, class_index):
        
        targets = data.targets
        self.exclude_target_index_list = [index for index in range(len(targets)) if targets[index] != class_index]
        self.indices = list(self.exclude_target_index_list)
    
    def __iter__(self, ):

        random.shuffle(self.indices)
        return iter(self.indices)
    
    def __len__(self, ):
        
        return len(self.indices)

class ClassSampler(Sampler):

    def __init__(self, data, class_index):
        targets = data.targets
        self.indices = list([index for index in range(len(targets)) if targets[index] == class_index])
    
    def __iter__(self):

        return iter(self.indices)

    def __len__(self, ):

        return len(self.indices)

class CrossSampler(Sampler):
    
    def __init__(self, data, class_index, total_num):
        """
        total_num: 决定参与生成基样本的总数目，算法以shuffle的方式随机采样
        """
        target = data.targets
        self.target_index_list = [index for index in range(len(target)) if target[index] == class_index]
        if total_num > len(self.target_index_list):
            raise Exception('the target class do not have so much samples')
        self.indices = list([self.target_index_list[index] for index in range(total_num)])
        

    def __iter__(self):
        
        t = 1000 * time.time()
        random.seed(int(t) % 2**32)
        random.shuffle(self.indices)
        print(self.indices)
        return iter(self.indices)

    def __len__(self, ):

        return len(self.indices)

class Dataer:
    def __init__(self, dataset_name='CIFAR-10', block_size=1):

        self.root_path = './data'
        self.dataset_name = dataset_name
        self.block_size = block_size

        # sample with value in [0, 1]
        transformer = transforms.Compose([transforms.ToTensor(), ])
        if dataset_name == 'MNIST':
            self.dataset = [
                torchvision.datasets.MNIST(root='./data/mnist', train=True, transform=transformer, download=True),
                torchvision.datasets.MNIST(root='./data/mnist', train=False, transform=transformer, download=True)                        
            ]   
        elif dataset_name == 'CIFAR-10':
            self.dataset = [
                torchvision.datasets.CIFAR10(root='./data/cifar-10-python', train=True, transform=transformer, download=True),
                torchvision.datasets.CIFAR10(root='./data/cifar-10-python', train=False, transform=transformer, download=True)                        
            ]   
        elif dataset_name == 'FMNIST':
            self.dataset = [
                torchvision.datasets.FashionMNIST(root='./data/fmnist', train=True, transform=transformer, download=True),
                torchvision.datasets.FashionMNIST(root='./data/fmnist', train=False, transform=transformer, download=True),
            ]
        elif dataset_name == 'ImageNet':
            root = '/home/jing/disk1/competition/robustness/dataset'
            # traindir = os.path.join(root, 'ILSVRC2012_img_train')
            three_part_dir = os.path.join(root, 'ILSVRC2012_img_val')
            
            traindir = os.path.join(self.root_path, 'TinyImageNet-10-python', 'train')
            valdir = os.path.join(self.root_path, 'TinyImageNet-10-python', 'test')
            
            self.dataset = [
                torchvision.datasets.ImageFolder(
                    traindir, 
                    transforms.Compose([
                        transforms.Resize(256),
                        transforms.CenterCrop(224),
                        transforms.ToTensor(),
                    ])                   
                ),
                torchvision.datasets.ImageFolder(
                    valdir,
                    transforms.Compose([
                        transforms.Resize(256),
                        transforms.CenterCrop(224),
                        transforms.ToTensor(),
                    ])
                ),
                torchvision.datasets.ImageFolder(
                    three_part_dir,
                    transforms.Compose([
                        transforms.Resize(256),
                        transforms.CenterCrop(224),
                        transforms.ToTensor(),
                    ])
                ),
            ]
        else:
            raise ValueError
        
        self.traindata_lenth = int(len(self.dataset[0]))
        self.testdata_lenth = int(len(self.dataset[1]))
        self.class_num = len(self.dataset[0].classes)

    
    def get_loader(self, batch_size=128, isTrain=True, isShuffle=True, isThree=False):

        if isTrain:
            return DataLoader(self.dataset[0], batch_size=batch_size, num_workers=4, shuffle=isShuffle, pin_memory=True)
        else:
            return DataLoader(self.dataset[1], batch_size=batch_size, num_workers=4, shuffle=isShuffle, pin_memory=True)
    
    def get_shuffle_loader(self, seed, batch_size=1, isTrain=False):

        if isTrain:
            shuffle_sampler = ShuffleSampler(self.dataset[0], seed=seed)
            return DataLoader(self.dataset[0], batch_size=batch_size, sampler=shuffle_sampler)
        else:
            shuffle_sampler = ShuffleSampler(self.dataset[1], seed=seed)
            return DataLoader(self.dataset[1], batch_size=batch_size, sampler=shuffle_sampler)

    def get_class_loader(self, class_index, isTrain=False, batch_size=8, isThreePart=False):
        """
        default batch size: 1
        return a loader, which give the sample with the label in class_index.
        """
        
        if isThreePart:
            class_sampler = ClassSampler(self.dataset[2], class_index)
            return DataLoader(self.dataset[2], batch_size=batch_size, sampler=class_sampler)

        if isTrain:
            class_sampler = ClassSampler(self.dataset[0], class_index)
            return DataLoader(self.dataset[0], batch_size=batch_size, sampler=class_sampler)
        else:
            class_sampler = ClassSampler(self.dataset[1], class_index)
            return DataLoader(self.dataset[1], batch_size=batch_size, sampler=class_sampler)
    
    def get_exclude_class_loader(self, class_index, isTrain=False, batch_size=100):

        if isTrain:
            exclude_class_sampler = ExculdeClassSampler(self.dataset[0], class_index)
            return DataLoader(self.dataset[0], batch_size=batch_size, sampler=exclude_class_sampler)
        else:
            exclude_class_sampler = ExculdeClassSampler(self.dataset[1], class_index)
            return DataLoader(self.dataset[1], batch_size=batch_size, sampler=exclude_class_sampler)

    def get_cross_class_loader(self, class_index, total_num, isTrain=False, batch_size=8):

        if isTrain:
            cross_sampler = CrossSampler(self.dataset[0], class_index, total_num)
            return DataLoader(self.dataset[0], batch_size=batch_size, sampler=cross_sampler)
        else:
            cross_sampler = CrossSampler(self.dataset[1], class_index, total_num)
            return DataLoader(self.dataset[1], batch_size=batch_size, sampler=cross_sampler)

    
    def get_sample(self, id=0, isShuffle=False, isTestData=True):
        """
        get one sample for test, default is 0-th image without random, where wo could chose the sample or random chose.
        """
        image, label = None, None
        if isShuffle:
            image, label = next(iter(self.get_loader(batch_size=1, isTrain=isTestData, isShuffle=True)))
        else:
            loader = self.get_loader(batch_size=1, isTrain=isTestData, isShuffle=False)
            for index, (sample, target) in enumerate(loader):
                if index == id:
                    image, label = sample, target
                    break
                else:
                    continue
        return image, label
    
    def get_shape(self, batch_size=1):

        loader = self.get_loader(batch_size=batch_size)
        samples, targets = next(iter(loader))

        ## ignore the batch size
        return samples.shape

    def get_reduction_shape(self, batch_size=1):
        
        loader = self.get_loader(batch_size=1)
        sample, _ = next(iter(loader))

        batch = batch_size
        channels = sample.shape[1]
        width = int(math.ceil(sample.shape[2] / self.block_size))
        height = int(math.ceil(sample.shape[3] / self.block_size))

        return (batch, channels, width, height)
        
    def get_channels(self, ):

        if self.dataset_name == 'CIFAR-10':
            return 3
        elif self.dataset_name == 'MNIST' or self.dataset_name == 'FMNIST':
            return 1
        else:
            return 3
    
    def get_dim(self, ):

        loader = self.get_loader(batch_size=1)
        sample, _ = next(iter(loader))

        if self.block_size == 1:
            return np.prod(sample.shape)
        else:
            ## Reduction strategy
            reduction_dim = np.prod(sample.shape[:2])
            for dim in sample.shape[2:]:
                reduction_dim *= int(math.ceil(dim / self.block_size))
            return reduction_dim


    def transfer(self, sample):

        dim = self.get_dim()        
        sample = torch.tensor(sample).to(dtype=torch.float)
        batch = int(np.prod(sample.shape) // dim)

        if self.block_size == 1:
            sample = sample.reshape(self.get_shape(batch_size=batch))
            return sample
        else:
            reduction_shape = self.get_reduction_shape(batch_size=batch)
            sample = sample.reshape(reduction_shape)

            ## expand
            expand_matrix = torch.repeat_interleave(sample, repeats=self.block_size, dim=2)
            expand_matrix = torch.repeat_interleave(expand_matrix, repeats=self.block_size, dim=3)
            
            ## cut out, if the expand matrix shape is greater than the origin size.
            origin_shape = self.get_shape()
            expand_matrix = expand_matrix[:, :, :origin_shape[2], :origin_shape[3]]
            return expand_matrix
    
    def retransfer(self, sample):

        """
        将真实的扰动降维为低维的扰动
        """
        
        return sample[:, :, ::self.block_size, ::self.block_size]
        
